import string

from entity import Entity
from model import Model
from atomsitecache import AtomsiteCache
from collections import OrderedDict
from pdboi.pdbents.exceptions import PDBentityNotFoundException

class Entry :
	'''this class represents a PDB entry.'''
	def __init__(s, dataconnector) :
		'''Entry needs a data connector to either mmcif or PDBML or PDB data'''
		s.dc = dataconnector
		s.entry_id = s.dc.getCategory("entry", "id", "first")
	def getId(s) : return (s.entry_id,)
	def getEntryId(s) : return s.entry_id
	def getDepositionDate(s) :
		if not hasattr(s,'depdate') :
			for fulldate in s.dc.getCategory("database_PDB_rev","date_original","iter") :
				if len(fulldate) > 3 :
					for yr in fulldate[0:4] : assert yr in "0123456789"
					s.depdate = fulldate ; break
		return s.depdate
	def getDepositors(s) :
		return list( s.dc.getCategory("audit_author","name","iter") )
	def getTitle(s) :
		'''return entry title, or None'''
		if not hasattr(s,'title') : s.title = s.dc.getCategory("struct","title","first")
		return s.title
	def prepareEntities(s) :
		'''prepare OrderedDict of entities'''
		if not hasattr(s,'entities') :
			s.entities = OrderedDict()
			for eid in s.dc.getCategory("entity","id","iter") : s.entities[eid] = Entity(s.dc,eid)
	def getEntity(s,eid) :
		'''return a list of Entity objects, make them first if required'''
		s.prepareEntities()
		return s.entities[eid]
	def getEntities(s) :
		'''return a list of Entity objects, make them first if required'''
		s.prepareEntities()
		return s.entities.values()
	def getModels(s) :
		'''return a list of model objects'''
		if not hasattr(s,'models') : s.makeModels()
		return s.models.values()
	def makeModels(s) :
		modelchains = AtomsiteCache.getCache(s.dc).getHierarchy()
		s.models = OrderedDict()
		for mid,chains in modelchains.items() : s.models[mid] = Model(s.dc,mid)
	def getModel(s, modelid) :
		if not hasattr(s,'models') : s.makeModels()
		try : return s.models[modelid]
		except KeyError : raise PDBentityNotFoundException("Model with id %s not found." % str(modelid))
	def getPDBaddressibleEntity(s, *args) :
		'''given one or more identifiers of model, chain, resnum, insertcode, resname, altcode, return corresponding model, chain, ModelledGroup, ModelledSubgroup.'''
		ent = s
		if len(args) > 0 : # model
			ent = ent.getModel(args[0])
		if len(args) > 1 : # chain
			ent = ent.getChain(args[1])
		if len(args) > 3 : # ModelledGroup
			ent = ent.getModelledGroup(args[2],args[3])
		if len(args) == 5 : # ModelledSubgroup - but only if unambiguous
			altcodes = ent.getAltcodesForGroupname(args[4])
			if len(altcodes) == 1 : args = tuple( args + tuple(altcodes) )
			elif len(altcodes) > 1 : raise PDBentityNotFoundException("More than 1 subgroup for "+str(args)+" : "+str(altcodes))
			else : raise PDBentityNotFoundException("No subgroup for "+str(args))
		if len(args) > 5 : # ModelledSubgroup
			ent = ent.getSubgroup(args[4],args[5])
		return ent
	def getNonpolymericNonwaterGroups(s) :
		'''returns subgroups belonging to nonpolymeric entities other than water'''
		if not hasattr(s,'models') : s.makeModels()
		ret = []
		for ent in s.getEntities() :
			if not ent.isType("non-polymer") or ent.isType("water") : continue
			for einst in ent.getInstances() :
				for CRIs in AtomsiteCache.getCache(s.dc).getESSmapCRI(ent.getEntityId(), einst.getStructAsymId()) :
					for chid,resnum,inscode in CRIs :
						for mid in s.models.keys() :
							ret.append( s.getPDBaddressibleEntity(mid, chid, resnum, inscode) )
		return ret
	def getCellDimensions(s) :
		'''cell dimensions from the cell table'''
		cellrow = s.dc.getCategory("cell",None,"first")
		cellvals = []
		for ckey in [ 'length_a', 'length_b', 'length_c', 'angle_alpha', 'angle_beta', 'angle_gamma', ] :
			cellvals.append( string.atof(cellrow[ckey]) )
		return cellvals
	def getSpacegroup(s) :
		'''spacegroup from the symmetry table'''
		return s.dc.getCategory("symmetry","space_group_name_H-M","first")
	def getResolutions(s) :
		'''return resolutions reported, the higher (numerically smaller) and lower respectively. None if not reported.'''
		'''spacegroup from the symmetry table'''
		resol, lowresol = None, None
		try : resol = s.dc.getCategory("refine","ls_d_res_high","first")
		except MMcifCategoryNotFoundException :
			try : resol = s.dc.getCategory("reflns","d_resolution_high","first")
			except MMcifCategoryNotFoundException : pass
		try : lowresol = s.dc.getCategory("refine","ls_d_res_low","first")
		except MMcifCategoryNotFoundException :
			try : lowresol = s.dc.getCategory("reflns","d_resolution_low","first")
			except MMcifCategoryNotFoundException : pass
		try : resol = string.atof(str(resol))
		except : pass
		try : lowresol = string.atof(str(lowresol))
		except : pass
		return resol, lowresol
	def getAtomElementCounts(s) :
		'''return a dictionary of element type and number of atoms modelled to that type'''
		return AtomsiteCache.getCache(s.dc).getElementCounts()
